import re
import numpy as np
import matplotlib.pyplot as plt

"""
画图函数
对训练过程得到的loss进行画图
"""

# 颜色配置
# colors = ['red','orange','green','cyan','blue','purple','brown','pink','magenta']
colors = ['red', 'orange', 'green', 'blue', 'purple', 'brown', 'pink', 'magenta']


def get_train_info(log_file):
    """读取文件的训练损失和验证损失值并保存"""
    train_losses = []
    val_losses = []
    train_f1s = []
    val_f1s = []
    with open(log_file, encoding="utf-8") as f:
        while True:
            lines = f.readline()
            if not lines:
                break
            if len(re.findall(r"(?<=train_loss:)\d+\.?\d*", lines)):
                train_loss = re.findall(r"(?<=train_loss:)\d+\.?\d*", lines)
                train_losses.append(float(train_loss[0]))
            if len(re.findall(r"(?<=val_loss:  )\d+\.?\d*", lines)):
                val_loss = re.findall(r"(?<=val_loss:  )\d+\.?\d*", lines)
                val_losses.append(float(val_loss[0]))
            if len(re.findall(r"(?<=train_f1:)\d+\.?\d*", lines)):
                train_f1 = re.findall(r"(?<=train_f1:)\d+\.?\d*", lines)
                train_f1s.append(float(train_f1[0]))
            if len(re.findall(r"(?<=val_f1:  )\d+\.?\d*", lines)):
                val_f1 = re.findall(r"(?<=val_f1:  )\d+\.?\d*", lines)
                val_f1s.append(float(val_f1[0]))
    return train_losses, val_losses, train_f1s, val_f1s


def plat_train_val(title, train_dict):
    """画单个文件的训练损失和验证损失"""
    plt.figure()
    train_loss = train_dict['train_loss']
    val_loss = train_dict['val_loss']
    train_f1 = train_dict['train_f1']
    val_f1 = train_dict['val_f1']
    epochs = len(train_loss)
    x = range(epochs)
    plt.plot(x, train_loss, label='train_loss', color=colors[0])
    plt.plot(x, val_loss, label='val_loss', color=colors[1])
    plt.plot(x, train_f1, label='train_f1', color=colors[2])
    plt.plot(x, val_f1, label='val_f1', color=colors[3])
    plt.title(title)
    plt.xlabel('epoch')
    plt.ylabel('loss')
    plt.legend(loc='best')
    plt.show()


if __name__ == '__main__':

    # GCN_log = "../logs/GCN/100k.log"
    GCN_log = "../logs/GCN/2025-04-05_13-32-26.log"
    # GAT_log = "../logs/GAT/100k.log"
    GAT_log = "../logs/GAT/2025-04-05_13-59-59.log"
    # GatedGCN_log = "../logs/GatedGCN/100k.log"
    GatedGCN_log = "../logs/GatedGCN/2024-09-27_09-56-25.log"

    info_dict = dict()
    info_dict['GCN'] = dict()
    info_dict['GAT'] = dict()
    info_dict['GatedGCN'] = dict()

    val_dict = dict()
    info_dict['GCN']['train_loss'], info_dict['GCN']['val_loss'],  \
    info_dict['GCN']['train_f1'], info_dict['GCN']['val_f1']= get_train_info(GCN_log)

    info_dict['GAT']['train_loss'], info_dict['GAT']['val_loss'],  \
    info_dict['GAT']['train_f1'], info_dict['GAT']['val_f1']= get_train_info(GAT_log)

    info_dict['GatedGCN']['train_loss'], info_dict['GatedGCN']['val_loss'],  \
    info_dict['GatedGCN']['train_f1'], info_dict['GatedGCN']['val_f1']= get_train_info(GatedGCN_log)

    plat_train_val('GCN', info_dict['GCN'])
    plat_train_val('GAT', info_dict['GAT'])
    plat_train_val('GatedGCN', info_dict['GatedGCN'])

